This is required for the deepseek-r1 model, where x is a sliced tensor that cannot be viewed.

we will do something else to avoid contiguous which will introduce mem copy here, all our quant should support non-contiguous tensor already... we hit any issue here?

Yes, I meet this issue for the chunked prefill path. see https://github.com/XiaobingSuper/ATOM/blob/xiaobing/oot_kimi/atom/plugin/attention_mla.py#L345.

then do contiguous at that place maybe, i don't like loss any perf

Yes, I updated the code, and do contiguous at plugin side.

self.num_heads * self.v_head_dim looks like not eaquals to 7168 for deepseek

This is o_proj's input, see atom path:
and plugin path:

The reason is that vllm do a_proj outside of the attention backend.

plugin path is also in our repo... then why we have to move o_proj out of attn

This change is for the fallback path, i.e., for plugin mode, but use vllm attn backend, because we use vllm MLAAttention class(here is self.attn), the forward path doesn't has o_proj, see https://github.com/vllm-project/vllm/blob/v0.15.1/vllm/attention/layer.py#L640, this is only for attention compute.

@valarLip The plugin mode using vllm atten backend will be like(setting ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1): https://github.com/zejunchen-zejun/ATOM/blob/zejun/plugin_for_atom_1223/recipes/vLLM-ATOM-OOT-Plugin-Backend.md#launching-server-of-vllm-with-atom-oot-plugin-platform

this is also comes from vllm? i would like we always use head_dim

Yes, vllm use head_size, see https://github.com/vllm-project/vllm/blob/v0.15.1/vllm/attention/layer.py#L579. Before this PR 6c40248, it also use head_size.

@zejunchen-zejun we need add this arg?

For vLLM side, it's calling path like https://github.com/vllm-project/vllm/blob/1892993bc18e243e2c05841314c5e9c06a80c70d/vllm/attention/layer.py#L675, it needs such a arg.

When use_mla is True, impl_cls is atom.model_ops.attention_mla.MLAAttention which now expects head_size (not head_dim). This code always includes head_dim in impl_args and then also adds head_size, so MLAAttention will receive an unexpected head_dim kwarg and raise at construction time. Build impl_args conditionally (only pass head_dim for MHA, and only pass head_size for MLA), or remove the unconditional head_dim entry before instantiating the MLA impl.